#!/usr/bin/env python
"""
数据导入脚本
从Excel文件导入时间序列数据到数据库
"""

import os
import sys
import django
import pandas as pd
import numpy as np
from datetime import datetime

# 设置Django环境
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'valar_hierarchical_forecasting.settings')
django.setup()

from grain.models import ForecastTask, TimeSeriesData, HierarchicalStructure


def get_region_mapping():
    """获取省份到区域的映射关系"""
    return {
        # 华北地区
        "北京市": "华北地区",
        "河北+天津": "华北地区", 
        "山西省": "华北地区",
        "内蒙古自治区": "华北地区",
        
        # 东北地区
        "辽宁省": "东北地区",
        "吉林省": "东北地区",
        "黑龙江省": "东北地区",
        
        # 华东地区
        "上海市": "华东地区",
        "江苏省": "华东地区",
        "浙江省": "华东地区",
        "安徽省": "华东地区",
        "福建省": "华东地区",
        "江西省": "华东地区",
        "山东省": "华东地区",
        
        # 华中地区
        "河南省": "华中地区",
        "湖北省": "华中地区",
        "湖南省": "华中地区",
        
        # 华南地区
        "广东+海南": "华南地区",
        "广西壮族自治区": "华南地区",
        
        # 西南地区
        "四川+重庆": "西南地区",
        "贵州省": "西南地区",
        "云南省": "西南地区",
        "西藏自治区": "西南地区",
        
        # 西北地区
        "陕西省": "西北地区",
        "甘肃省": "西北地区",
        "青海省": "西北地区",
        "宁夏回族自治区": "西北地区",
        "新疆维吾尔自治区": "西北地区"
    }


def create_hierarchical_structure():
    """创建默认的层级结构配置"""
    try:
        # 检查是否已存在默认配置
        if HierarchicalStructure.objects.filter(is_default=True).exists():
            print("✓ 默认层级结构已存在")
            return HierarchicalStructure.objects.filter(is_default=True).first()
        
        # 创建默认的中国区域层级结构
        default_hierarchy = {
            "regions": {
                "华北地区": ["北京市", "河北+天津", "山西省", "内蒙古自治区"],
                "东北地区": ["辽宁省", "吉林省", "黑龙江省"],
                "华东地区": ["上海市", "江苏省", "浙江省", "安徽省", "福建省", "江西省", "山东省"],
                "华中地区": ["河南省", "湖北省", "湖南省"],
                "华南地区": ["广东+海南", "广西壮族自治区"],
                "西南地区": ["四川+重庆", "贵州省", "云南省", "西藏自治区"],
                "西北地区": ["陕西省", "甘肃省", "青海省", "宁夏回族自治区", "新疆维吾尔自治区"]
            },
            "levels": ["national", "regional", "provincial"],
            "description": "中国七大地理区域划分"
        }
        
        hierarchy = HierarchicalStructure.objects.create(
            name="中国区域层级结构",
            description="基于中国七大地理区域的层级结构配置",
            hierarchy_config=default_hierarchy,
            is_default=True
        )
        
        print("✓ 默认层级结构创建成功")
        return hierarchy
        
    except Exception as e:
        print(f"✗ 创建层级结构失败: {str(e)}")
        return None


def load_excel_data(file_path):
    """加载Excel数据"""
    try:
        print(f"正在读取Excel文件: {file_path}")
        
        # 读取Excel文件
        df = pd.read_excel(file_path, sheet_name=0)
        print(f"✓ Excel文件读取成功，数据形状: {df.shape}")
        
        # 显示数据的前几行和列名
        print(f"✓ 列名: {list(df.columns)}")
        print(f"✓ 前5行数据:")
        print(df.head())
        
        return df
        
    except Exception as e:
        print(f"✗ Excel文件读取失败: {str(e)}")
        return None


def process_and_import_data(df, task_name="粮食产量预测任务"):
    """处理并导入数据到数据库"""
    try:
        print("开始处理和导入数据...")
        
        # 1. 创建预测任务
        print("1. 创建预测任务...")
        task = ForecastTask.objects.create(
            task_name=task_name,
            description="从Excel文件导入的粮食产量时间序列数据",
            forecast_horizon=5,
            start_year=int(df.iloc[:, 0].min()) if len(df) > 0 else None,
            end_year=int(df.iloc[:, 0].max()) if len(df) > 0 else None,
            status='pending'
        )
        print(f"✓ 预测任务创建成功，任务ID: {task.id}")
        
        # 2. 处理数据格式
        print("2. 处理数据格式...")
        
        # 找到时间列
        time_column = None
        for col in df.columns:
            if '时间' in str(col) or '年份' in str(col) or '年' in str(col):
                time_column = col
                break
        
        if time_column is None:
            # 假设第一列是时间列
            time_column = df.columns[0]
        
        print(f"✓ 时间列: {time_column}")
        
        # 设置时间列为索引
        df_processed = df.set_index(time_column)
        
        # 获取区域映射
        region_mapping = get_region_mapping()
        
        # 3. 转换为长格式并导入数据
        print("3. 转换数据格式并导入数据库...")
        
        data_objects = []
        total_records = 0
        
        for year in df_processed.index:
            for province in df_processed.columns:
                value = df_processed.loc[year, province]
                
                # 跳过空值
                if pd.isna(value) or value == 0:
                    continue
                
                # 获取区域
                region = region_mapping.get(province, '未知区域')
                
                # 创建数据对象
                data_objects.append(TimeSeriesData(
                    task=task,
                    year=int(year),
                    province=province,
                    region=region,
                    value=float(value),
                    data_type='original'
                ))
                
                total_records += 1
                
                # 批量插入（每1000条记录插入一次）
                if len(data_objects) >= 1000:
                    TimeSeriesData.objects.bulk_create(data_objects, ignore_conflicts=True)
                    print(f"  已导入 {len(data_objects)} 条记录...")
                    data_objects = []
        
        # 插入剩余的记录
        if data_objects:
            TimeSeriesData.objects.bulk_create(data_objects, ignore_conflicts=True)
            print(f"  已导入 {len(data_objects)} 条记录...")
        
        print(f"✓ 数据导入完成，总计导入 {total_records} 条记录")
        
        # 4. 更新任务信息
        task.status = 'completed'
        task.save()
        
        # 5. 显示导入统计
        print("\n数据导入统计:")
        print(f"- 任务ID: {task.id}")
        print(f"- 任务名称: {task.task_name}")
        print(f"- 数据年份范围: {task.start_year} - {task.end_year}")
        print(f"- 总记录数: {total_records}")
        
        # 按省份统计
        province_stats = TimeSeriesData.objects.filter(task=task).values('province').distinct().count()
        print(f"- 省份数量: {province_stats}")
        
        # 按区域统计
        region_stats = TimeSeriesData.objects.filter(task=task).values('region').distinct().count()
        print(f"- 区域数量: {region_stats}")
        
        return task
        
    except Exception as e:
        print(f"✗ 数据处理和导入失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return None


def main():
    """主函数"""
    print("=" * 60)
    print("分层时间序列数据导入脚本")
    print("=" * 60)
    
    # 1. 创建层级结构
    print("\n1. 创建层级结构...")
    hierarchy = create_hierarchical_structure()
    if not hierarchy:
        print("层级结构创建失败，退出程序")
        return
    
    # 2. 读取Excel文件
    excel_file_path = "/Users/jm/Desktop/分层时间序列材料/数据/原始数据/全部数据--合并.xlsx"
    
    print(f"\n2. 读取Excel文件...")
    df = load_excel_data(excel_file_path)
    if df is None:
        print("Excel文件读取失败，退出程序")
        return
    
    # 3. 处理并导入数据
    print(f"\n3. 处理并导入数据...")
    task = process_and_import_data(df, "粮食产量预测任务")
    if task is None:
        print("数据导入失败，退出程序")
        return
    
    print("\n" + "=" * 60)
    print("数据导入完成！")
    print(f"任务ID: {task.id}")
    print("可以使用以下API接口进行预测:")
    print("- POST /forecast/bottom-up  (自底向上预测)")
    print("- POST /forecast/mint       (MinT协调预测)")
    print("- POST /forecast/ilr-var    (ILR-VAR预测)")
    print("=" * 60)


if __name__ == "__main__":
    main()

